import numpy as np 

# A is nxd
# B is mxd
def allpairwise_cosim(A, B):
    # cij = cos(ai, bj)
    d = A.shape[1]
    # convert matrix size to n x m x d 
    A_re = A.reshape((-1, 1, d))
    B_re = B.reshape((1, -1, d))
    dot_product = A_re * B_re
    norm_product = (np.sqrt((A_re * A_re).sum(axis=2)) * 
        np.sqrt((B_re * B_re).sum(axis=2)))
    dot_product = dot_product.sum(axis=2)
    C = dot_product / norm_product
    return C

def onepair_cosim(a, b): 
    cos_sim = np.dot(a, b)/(np.linalg.norm(a)*np.linalg.norm(b))
    return cos_sim

if __name__ == "__main__":
    A = np.array([[1, 2], [2, -1], [3, -1]], dtype=float)
    B = np.array([[2, 3], [4, 5], [3, -5], [4, 6]], dtype=float)
    C = allpairwise_cosim(A, B)
    assert C.shape == (len(A), len(B))
    for i in range(len(A)): 
        for j in range(len(B)): 
            assert np.allclose(C[i][j], onepair_cosim(A[i], B[j]))
